# doris_to_mysql_insert.py

import pymysql
import pandas as pd
from config import DORIS_CONFIG, MYSQL_CONFIG

def infer_sql_type(dtype):
    """根据 pandas 类型推断 MySQL 字段类型"""
    if pd.api.types.is_integer_dtype(dtype):
        return "BIGINT"
    elif pd.api.types.is_float_dtype(dtype):
        return "DOUBLE"
    elif pd.api.types.is_datetime64_any_dtype(dtype):
        return "DATETIME"
    else:
        return "TEXT"


def create_table_from_df(cursor, table_name, df):
    """根据 DataFrame 自动生成建表语句"""
    columns = []
    for col_name, dtype in df.dtypes.items():
        sql_type = infer_sql_type(dtype)
        columns.append(f"`{col_name}` {sql_type}")

    col_definitions = ", ".join(columns)
    create_table_sql = f"CREATE TABLE IF NOT EXISTS `{table_name}` ({col_definitions});"

    try:
        cursor.execute(create_table_sql)
        print(f"表 `{table_name}` 已创建或已存在")
    except Exception as e:
        print(f"创建表失败: {e}")


def insert_data_from_df(cursor, table_name, df):
    """从 DataFrame 批量插入数据"""
    placeholders = ", ".join(["%s"] * len(df.columns))
    columns = ", ".join([f"`{col}`" for col in df.columns])
    insert_sql = f"INSERT INTO `{table_name}` ({columns}) VALUES ({placeholders});"

    data_tuples = list(df.itertuples(index=False, name=None))

    try:
        cursor.executemany(insert_sql, data_tuples)
        print(f"成功插入 {len(data_tuples)} 条数据到 `{table_name}`")
    except Exception as e:
        print(f"插入数据失败: {e}")


def main():
    # 连接 Doris 数据库
    try:
        doris_conn = pymysql.connect(**DORIS_CONFIG)
        print("正在连接 Doris 数据库...")
        with doris_conn.cursor() as doris_cursor:
            query = "SELECT * FROM transport_task_final LIMIT 25"
            doris_cursor.execute(query)
            result = doris_cursor.fetchall()
            df = pd.DataFrame(result)
        print("成功从 Doris 查询到 25 条数据")
    except Exception as e:
        print(f"连接 Doris 或查询失败: {e}")
        return
    finally:
        doris_conn.close()
        print("已关闭 Doris 数据库连接")

    if df.empty:
        print("没有查询到任何数据，请检查 SQL 或表是否存在")
        return

    # 连接 MySQL 数据库
    try:
        mysql_conn = pymysql.connect(**MYSQL_CONFIG)
        print("正在连接 MySQL 数据库...")
        with mysql_conn.cursor() as mysql_cursor:
            create_table_from_df(mysql_cursor, "transport_task_final", df)
            insert_data_from_df(mysql_cursor, "transport_task_final", df)
        mysql_conn.commit()
        print("事务提交完成")
    except Exception as e:
        print(f"操作 MySQL 失败: {e}")
    finally:
        mysql_conn.close()
        print("已关闭 MySQL 数据库连接")


if __name__ == "__main__":
    main()